解題程式碼

var distanceK = function (root, target, k) {
  const graphMap = new Map();
  const visited = new Set();
  const res = [];
  const queue = [{ nodeVal: target.val, layer: 0 }];

  const buildGraph = (node, parent) => {
    if (!node) return;
    const neighbor = [];

    if (parent) neighbor.push(parent.val);
    if (node.left) {
      neighbor.push(node.left.val);
      buildGraph(node.left, node);
    }
    if (node.right) {
      neighbor.push(node.right.val);
      buildGraph(node.right, node);
    }

    graphMap.set(node.val, neighbor);
  };
  buildGraph(root, null);

  while (queue.length) {
    const { nodeVal, layer } = queue.shift();
    if (visited.has(nodeVal)) continue;

    visited.add(nodeVal);
    if (layer === k) res.push(nodeVal);
    if (layer > k) break;

    for (const ele of graphMap.get(nodeVal)) {
      queue.push({ nodeVal: ele, layer: layer + 1 });
    }
  }

  return res;
};

解題思路、演算法

這題可以將樹轉成圖的形式，將圖的資料存在 hashMap 裡面，以下面的 tree 為範例的話，就是:

{
  3: [5, 1],
  5: [3, 6, 2],
  6: [5],
  2: [5, 7, 4],
  7: [2],
  4: [2],
  1: [3, 0, 8],
  0: [1],
  8: [1]
}

建立好圖後，將 target 節點當作圖的出發點，用 queue 儲存，然後每次從 queue 裡取出節點並透過 hashMap 查表，就能找出鄰近節點，經過 k 路徑所到的節點就都會是題目要的結果。

Tree is special form of graph.

解法的時間、空間複雜度

時間複雜度: O(m + n)，m = edges，為 n - 1，n = nodes in the tree
空間複雜度: O(n)

參考資料

All Nodes Distance K In A Binary Tree - Performing Bidirectional Search On A Tree Using A Hashtable

leetcode 中文 | All Nodes Distance K in Binary Tree | Amazon 亞馬遜考題 | Leetcode 863 | Python